
#include "PostProcessingPipeline.h"
#include "D3D11Client.h"

PSPostProcessingPipeline::PSPostProcessingPipeline(D3D11Client * d3d, CComPtrEx<ID3D11Texture2D>& inputTexture)
	: PostProcessingPipeline(d3d, inputTexture)
{
	m_width = cfg->cWidth;
	m_height = cfg->cHeight;
	m_adaptedLuminance = m_currentLuminance = 1.0f;
	for(auto i = 0; i < LUMINANCE_HISTORY_LENGTH; i++)
	{
		m_luminanceHistory[i] = m_currentLuminance;
	}
	m_currentBuffer = 0;
}

bool PSPostProcessingPipeline::Init(CComPtrEx<ID3D11Texture2D>& outputTexture)
{
	HRESULT hr;
	auto device = Dev;

	D3D11_TEXTURE2D_DESC tDesc;
	D3D11_SHADER_RESOURCE_VIEW_DESC srvDesc;
	D3D11_RENDER_TARGET_VIEW_DESC rtvDesc;

	ZeroMemory(&tDesc, sizeof(D3D11_TEXTURE2D_DESC));
	tDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_RENDER_TARGET;
	tDesc.Format = DXGI_FORMAT_R16G16B16A16_FLOAT;
	tDesc.Width = m_width;
	tDesc.Height = m_height;
	tDesc.ArraySize = 1;
	tDesc.MipLevels = 1;
	tDesc.SampleDesc.Count = 1;
	tDesc.Usage = D3D11_USAGE_DEFAULT;
	hr = device->CreateTexture2D(&tDesc, nullptr, &m_outputTexture);
	if (FAILED(hr))
		return false;
	outputTexture = m_outputTexture;

	hr = device->CreateTexture2D(&tDesc, nullptr, &m_bloomTexture);
	if (FAILED(hr))
		return false;

	hr = device->CreateTexture2D(&tDesc, nullptr, &m_gaussianBlurTempTexture);
	if (FAILED(hr))
		return false;

	hr = device->CreateTexture2D(&tDesc, nullptr, &m_lightStreaksTexture);
	if (FAILED(hr))
		return false;

	hr = device->CreateTexture2D(&tDesc, nullptr, &m_lightStreaksTempTexture);
	if (FAILED(hr))
		return false;
	
	ZeroMemory(&srvDesc, sizeof(D3D11_SHADER_RESOURCE_VIEW_DESC));
	srvDesc.Format = DXGI_FORMAT_R16G16B16A16_FLOAT;
	srvDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
	srvDesc.Texture2D.MipLevels = 1;
	hr = device->CreateShaderResourceView(m_inputTexture, &srvDesc, &m_inputTextureSRV);
	if (FAILED(hr))
		return false;

	hr = device->CreateShaderResourceView(m_bloomTexture, &srvDesc, &m_bloomTextureSRV);
	if (FAILED(hr))
		return false;

	hr = device->CreateShaderResourceView(m_gaussianBlurTempTexture, &srvDesc, &m_gaussianBlurTempTextureSRV);
	if (FAILED(hr))
		return false;

	hr = device->CreateShaderResourceView(m_lightStreaksTexture, &srvDesc, &m_lightStreaksTextureSRV);
	if (FAILED(hr))
		return false;

	hr = device->CreateShaderResourceView(m_lightStreaksTempTexture, &srvDesc, &m_lightStreaksTempTextureSRV);
	if (FAILED(hr))
		return false;

	ZeroMemory(&rtvDesc, sizeof(D3D11_RENDER_TARGET_VIEW_DESC));
	rtvDesc.Format = DXGI_FORMAT_R16G16B16A16_FLOAT;
	rtvDesc.ViewDimension = D3D11_RTV_DIMENSION_TEXTURE2D;
	rtvDesc.Texture2D.MipSlice = 0;
	hr = device->CreateRenderTargetView(m_bloomTexture, &rtvDesc, &m_bloomTextureRTV);
	if (FAILED(hr))
		return false;
	hr = device->CreateRenderTargetView(m_outputTexture, &rtvDesc, &m_outputTextureRTV);
	if (FAILED(hr))
		return false;

	hr = device->CreateRenderTargetView(m_gaussianBlurTempTexture, &rtvDesc, &m_gaussianBlurTempTextureRTV);
	if (FAILED(hr))
		return false;

	hr = device->CreateRenderTargetView(m_lightStreaksTexture, &rtvDesc, &m_lightStreaksTextureRTV);
	if (FAILED(hr))
		return false;

	hr = device->CreateRenderTargetView(m_lightStreaksTempTexture, &rtvDesc, &m_lightStreaksTempTextureRTV);
	if (FAILED(hr))
		return false;

	D3D11_BUFFER_DESC bDesc;
	ZeroMemory(&bDesc, sizeof(D3D11_BUFFER_DESC));
	bDesc.BindFlags = D3D11_BIND_VERTEX_BUFFER;
	bDesc.Usage = D3D11_USAGE_DEFAULT;
	bDesc.StructureByteStride = sizeof(D3DXVECTOR4);
	bDesc.ByteWidth = 6 * sizeof(D3DXVECTOR4);

	D3DXVECTOR4 vertexData [] =
	{
		D3DXVECTOR4(-1, -1, 0, 1),
		D3DXVECTOR4(-1,  1, 0, 1),
		D3DXVECTOR4( 1,  1, 0, 1),

		D3DXVECTOR4(-1, -1, 0, 1),
		D3DXVECTOR4( 1,  1, 0, 1),
		D3DXVECTOR4( 1, -1, 0, 1),
	};

	D3D11_SUBRESOURCE_DATA initData;
	ZeroMemory(&initData, sizeof(D3D11_SUBRESOURCE_DATA));
	initData.pSysMem = vertexData;

	hr = device->CreateBuffer(&bDesc, &initData, &m_quadBuffer);
	if (FAILED(hr))
		return false;

	ZeroMemory(&bDesc, sizeof(D3D11_BUFFER_DESC));
	bDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	bDesc.Usage = D3D11_USAGE_DEFAULT;
	bDesc.ByteWidth = sizeof(csCopyTextureBuffer);
	hr = device->CreateBuffer(&bDesc, nullptr, &m_csCopyTextureBuffer);
	if (FAILED(hr))
		return false;

	D3D11_INPUT_ELEMENT_DESC ldesc[ ] = 
	{
		{ "SV_POSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
	};

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\PassThroughVS.fx", "VS", m_passThroughVS))
		return false;

	if (!GetClient()->CreateInputLayout(ldesc, m_passThroughVS, m_layout))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_CalculateLuminance.fx", "CalculateLuminance", m_calculateLuminancePS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_Reduce.fx", "Reduce", m_reducePS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_CalculateBloomTexture.fx", "CalculateBloomTexture", m_calculateBloomTexturePS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_CalculateLightStreaksTexture.fx", "CalculateLightStreaksTexture", m_calculateLightStreaksTexturePS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_CreateLightStreaks.fx", "CreateLightStreaks", m_createLightStreaksTexturePS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_GaussianBlur.fx", "BloomTextureBlurVerticalPass", m_gaussianBlurVerticalPassPS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_GaussianBlur.fx", "BloomTextureBlurHorizontalPass", m_gaussianBlurHorizontalPassPS))
		return false;

	if (!GET_SHADER(GetClient(), "Modules\\D3D11Shaders\\PSPipeline\\HDR_PS_CopyTextureWithExposure.fx", "CopyTextureWithExposure", m_copyTextureWithExposurePS))
		return false;
	
	int currSizeX = m_width;
	int currSizeY = m_height;

	while ((currSizeX > 1) || (currSizeY > 1))
	{
		CalcLuminance curr;
		currSizeX = (int)ceil((float)currSizeX / 8);
		currSizeY = (int)ceil((float)currSizeY / 8);
		curr.width = currSizeX;
		curr.height = currSizeY;

		ZeroMemory(&tDesc, sizeof(D3D11_TEXTURE2D_DESC));
		tDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_RENDER_TARGET;
		tDesc.Format = DXGI_FORMAT_R32_FLOAT;
		tDesc.Width = curr.width;
		tDesc.Height = curr.height;
		tDesc.ArraySize = 1;
		tDesc.MipLevels = 1;
		tDesc.SampleDesc.Count = 1;
		tDesc.Usage = D3D11_USAGE_DEFAULT;
		hr = device->CreateTexture2D(&tDesc, nullptr, &curr.texture);
		if (FAILED(hr))
			return false;

		ZeroMemory(&srvDesc, sizeof(D3D11_SHADER_RESOURCE_VIEW_DESC));
		srvDesc.Format = DXGI_FORMAT_R32_FLOAT;
		srvDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
		srvDesc.Texture2D.MipLevels = 1;
		hr = device->CreateShaderResourceView(curr.texture, &srvDesc, &curr.srv);
		if (FAILED(hr))
			return false;

		ZeroMemory(&rtvDesc, sizeof(D3D11_RENDER_TARGET_VIEW_DESC));
		rtvDesc.Format = DXGI_FORMAT_R32_FLOAT;
		rtvDesc.ViewDimension = D3D11_RTV_DIMENSION_TEXTURE2D;
		rtvDesc.Texture2D.MipSlice = 0;
		hr = device->CreateRenderTargetView(curr.texture, &rtvDesc, &curr.rtv);
		if (FAILED(hr))
			return false;

		m_calculateLuminanceData.push_back(curr);
	}

	ZeroMemory(&tDesc, sizeof(D3D11_TEXTURE2D_DESC));
	tDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	tDesc.Format = DXGI_FORMAT_R32_FLOAT;
	tDesc.Width = 2;
	tDesc.Height = 1;
	tDesc.ArraySize = 1;
	tDesc.MipLevels = 0;
	tDesc.SampleDesc.Count = 1;
	tDesc.Usage = D3D11_USAGE_DEFAULT;
	hr = device->CreateTexture2D(&tDesc, nullptr, &m_luminanceTexture);
	if (FAILED(hr))
		return false;

	ZeroMemory(&srvDesc, sizeof(D3D11_SHADER_RESOURCE_VIEW_DESC));
	srvDesc.Format = DXGI_FORMAT_R32_FLOAT;
	srvDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
	srvDesc.Texture2D.MipLevels = 1;
	hr = device->CreateShaderResourceView(m_luminanceTexture, &srvDesc, &m_luminanceTextureSRV);
	if (FAILED(hr))
		return false;
	
	tDesc.Width = 1;
	tDesc.BindFlags = 0;
	tDesc.Usage = D3D11_USAGE_STAGING;
	tDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	for(auto i = 0; i < 3; i++)
	{
		hr = device->CreateTexture2D(&tDesc, nullptr, &m_luminanceTextureRW[i]);
		if (FAILED(hr))
			return false;
	}

	D3D11_DEPTH_STENCIL_DESC dsDesc;
	ZeroMemory(&dsDesc, sizeof(D3D11_DEPTH_STENCIL_DESC));
	dsDesc.DepthEnable = false;
	//m_dsState

	return true;
}

inline float CalcMiddleGray(float averageLuminance)
{
	return (1.03f - 2 / (2 + log10(averageLuminance + 1)));
}

void PSPostProcessingPipeline::RunPipeline(double interval, bool isPause)
{
	auto context = iCtx;
	ID3D11ShaderResourceView * emptySRVs[] = { nullptr, nullptr, nullptr };
	ID3D11RenderTargetView * emptyRTVs = { nullptr };

	float colorMask[] = { 1, 1, 1, 1 };

	context->OMSetBlendState(BS_NoBlend, colorMask, 0xffffffff);

	csCopyTextureBuffer args;
	ZeroMemory(&args, sizeof(csCopyTextureBuffer));
	
	context->PSSetConstantBuffers(0, 1, &m_csCopyTextureBuffer);


	CalcLuminance& curr = m_calculateLuminanceData[0];
	D3D11_VIEWPORT vp = { 0 };
	vp.Width = (float)curr.width;
	vp.Height = (float)curr.height;
	context->OMSetRenderTargets(1, &curr.rtv, nullptr);
	context->RSSetViewports(1, &vp);
	float color[] = { 0.0f, 0.0f, 0.0f, 0.0f };
	context->ClearRenderTargetView(curr.rtv, color);
	context->IASetInputLayout(m_layout);
	context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
	UINT strides[] = { sizeof(D3DXVECTOR4) };
	UINT offsets[] = { 0 };
	context->IASetVertexBuffers(0, 1, &m_quadBuffer, strides, offsets);

	context->VSSetShader(m_passThroughVS, nullptr, 0);

	args.textureSize[0] = m_width;
	args.textureSize[1] = m_height;

	args.secondTextureSize[0] = curr.width;
	args.secondTextureSize[1] = curr.height;
	context->UpdateSubresource(m_csCopyTextureBuffer, 0, nullptr, &args, 0, 0);

	context->PSSetShader(m_calculateLuminancePS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_inputTextureSRV);
	context->Draw(6, 0);

	context->OMSetRenderTargets(1, &emptyRTVs, nullptr);
	
	context->PSSetShader(m_reducePS, nullptr, 0);

	auto dataSize = m_calculateLuminanceData.size();
	for(UINT i =  1; i < dataSize; i++)
	{
		CalcLuminance& curr = m_calculateLuminanceData[i];

		vp.Width = (float)curr.width;
		vp.Height = (float)curr.height;

		args.textureSize[0] = m_calculateLuminanceData[i - 1].width;
		args.textureSize[1] = m_calculateLuminanceData[i - 1].height;

		args.secondTextureSize[0] = curr.width;
		args.secondTextureSize[1] = curr.height;
		context->UpdateSubresource(m_csCopyTextureBuffer, 0, nullptr, &args, 0, 0);

		context->RSSetViewports(1, &vp);
		context->OMSetRenderTargets(1, &curr.rtv, nullptr);
		context->ClearRenderTargetView(curr.rtv, color);
		context->PSSetShaderResources(0, 1, &m_calculateLuminanceData[i - 1].srv);
		context->Draw(6, 0);
	}

	D3D11_BOX copyBox = { 0 };
	copyBox.right = 1;
	copyBox.bottom = 1;
	copyBox.back = 1;
	context->CopySubresourceRegion(m_luminanceTextureRW[m_currentBuffer], 0, 0, 0, 0, m_calculateLuminanceData[dataSize - 1].texture, 0, &copyBox);

	D3D11_MAPPED_SUBRESOURCE subResource;
	float luminanceBuffer[2];

	auto bufferToRead = m_currentBuffer + 1;
	if (bufferToRead > 2)
	{
		bufferToRead -= 3;
	}

	context->Map(m_luminanceTextureRW[bufferToRead], 0, D3D11_MAP_READ, 0, &subResource);
	auto pData = (float *)subResource.pData;
	auto luminance = *pData;
	context->Unmap(m_luminanceTextureRW[bufferToRead], 0);

	m_currentBuffer++;
	if (m_currentBuffer > 2)
	{
		m_currentBuffer -= 3;
	}

	CLAMP_VALUE(luminance, 0.7f, 7.0f);
	CLAMP_VALUE(m_currentLuminance, 0.7f, 7.0f);

	static const float tau = 0.3f;
	if (!isPause)
	{
		auto timeStep = interval;
		auto lumDiff = m_currentLuminance - luminance;
		int total = 0;
		for(UINT i = 0; i < LUMINANCE_HISTORY_LENGTH; i++)
		{
			if (m_luminanceHistory[i] > luminance)
			{
				total += 1;
			}
		}
		if (total == 0 || total == 16)
		{
			luminance = m_currentLuminance + (luminance - m_currentLuminance) * (1 - exp(-timeStep * tau));
		}
		else
		{
			luminance = m_currentLuminance;
		}
		m_currentLuminance = luminance;
	}
	else
	{
		luminance = m_currentLuminance;
	}

	memcpy(m_luminanceHistory, &m_luminanceHistory[1], sizeof(float) * (LUMINANCE_HISTORY_LENGTH - 1));
	m_luminanceHistory[LUMINANCE_HISTORY_LENGTH - 1] = luminance;
	
	auto middleGray = CalcMiddleGray(luminance) * 2.0f;
	luminanceBuffer[0] = luminance;
	luminanceBuffer[1] = middleGray;
	context->UpdateSubresource(m_luminanceTexture, 0, nullptr, luminanceBuffer, 0, 0);

	args.textureSize[0] = m_width;
	args.textureSize[1] = m_height;
	context->UpdateSubresource(m_csCopyTextureBuffer, 0, nullptr, &args, 0, 0);

	vp.Width = (float)m_width;
	vp.Height = (float)m_height;
	context->OMSetRenderTargets(1, &m_bloomTextureRTV, nullptr);
	context->RSSetViewports(1, &vp);

	context->PSSetShader(m_calculateBloomTexturePS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_inputTextureSRV);
	context->PSSetShaderResources(1, 1, &m_luminanceTextureSRV);

	context->Draw(6, 0);

	//light streaks

	/*context->OMSetRenderTargets(1, &m_lightStreaksTextureRTV, nullptr);

	context->PSSetShader(m_calculateLightStreaksTexturePS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_inputTextureSRV);
	context->PSSetShaderResources(1, 1, &m_luminanceTextureSRV);

	context->Draw(6, 0);

	context->PSSetShaderResources(0, 3, emptySRVs);

	context->OMSetRenderTargets(1, &m_lightStreaksTempTextureRTV, nullptr);

	context->PSSetShader(m_createLightStreaksTexturePS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_lightStreaksTextureSRV);

	context->Draw(6, 0);*/

	//gaussian blur, vertical pass

	context->OMSetRenderTargets(1, &m_gaussianBlurTempTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurVerticalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_bloomTextureSRV);

	context->Draw(6, 0);
	
	//gaussian blur, horizontal pass

	context->PSSetShaderResources(0, 3, emptySRVs);

	context->OMSetRenderTargets(1, &m_bloomTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurHorizontalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_gaussianBlurTempTextureSRV);

	context->Draw(6, 0);

	//gaussian blur, vertical pass
	/*context->PSSetShaderResources(0, 1, emptySRVs);

	context->OMSetRenderTargets(1, &m_gaussianBlurTempTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurVerticalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_lightStreaksTempTextureSRV);

	context->Draw(6, 0);
	
	//gaussian blur, horizontal pass

	context->PSSetShaderResources(0, 1, emptySRVs);

	context->OMSetRenderTargets(1, &m_lightStreaksTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurHorizontalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_gaussianBlurTempTextureSRV);

	context->Draw(6, 0);*/

	//gaussian blur #2, vertical pass
	
	/*context->PSSetShaderResources(0, 1, emptySRVs);

	context->OMSetRenderTargets(1, &m_gaussianBlurTempTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurVerticalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_bloomTextureSRV);

	context->Draw(6, 0);
	
	//gaussian blur, horizontal pass

	context->PSSetShaderResources(0, 1, emptySRVs);

	context->OMSetRenderTargets(1, &m_bloomTextureRTV, nullptr);
	context->PSSetShader(m_gaussianBlurHorizontalPassPS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_gaussianBlurTempTextureSRV);

	context->Draw(6, 0);*/

	//final pass
	context->OMSetRenderTargets(1, &m_outputTextureRTV, nullptr);
	context->PSSetShader(m_copyTextureWithExposurePS, nullptr, 0);
	context->PSSetShaderResources(0, 1, &m_inputTextureSRV);
	context->PSSetShaderResources(1, 1, &m_bloomTextureSRV);
	//context->PSSetShaderResources(2, 1, &m_lightStreaksTempTextureSRV);
	context->PSSetShaderResources(3, 1, &m_luminanceTextureSRV);

	context->Draw(6, 0);

	//context->CopyResource(m_outputTexture, m_inputTexture);
}